from socket import *
import struct

#UDPsocket = socket(AF_INET,SOCK_DGRAM)
UDPsocket = None
def downLoad():
    ip = "192.168.31.108"
    sendAddr = (ip, 69)
    cmd_buf = struct.pack(b"!H8sb5sb", 1, b"test.jpg", 0, b"octet", 0)
    UDPsocket.sendto(cmd_buf, sendAddr)
    p_num = 0
    recvFile = ""
    while True:
        recvData, recvAddr = UDPsocket.recvfrom(1024)
        recvDataLen = len(recvData)
        cmd_Tuple = struct.unpack("!HH", recvData[:4])
        cmd = cmd_Tuple[0]
        currentPackNum = cmd_Tuple[1]
        if cmd == 3:
            if currentPackNum == 1:
                recvFile = open("test.jpg", 'wb')
            if currentPackNum == (p_num + 1):
                recvFile.write(recvData[4:])
                p_num += 1
                print("接收到第%d块数据" % p_num)

                ackData = struct.pack("!HH", 4, p_num)
                UDPsocket.sendto(ackData, recvAddr)

            if recvDataLen < 516:
                recvFile.close()
                print("文件已经成功下载")
                break
        elif cmd == 5:
            print("错误应答码：%d"%currentPackNum)
            break

def upLoad():
    ip = "192.168.31.108"
    sendAddr = (ip, 69)
    cmd_buf = struct.pack(b"!H8sb5sb", 2, b"test.jpg", 0, b"octet", 0)
    UDPsocket.sendto(cmd_buf, sendAddr)
    p_num = 0
    readFile = ""
    finishNum = -1
    while True:
        recvData, recvAddr = UDPsocket.recvfrom(1024)
        cmd_Tuple = struct.unpack("!HH", recvData[:4])
        cmd = cmd_Tuple[0]
        currentPackNum = cmd_Tuple[1]
        if finishNum == currentPackNum:#收到最后一个ack编号，退出
            readFile.close()
            print("文件已经成功发送")
            break
        if cmd == 4 and currentPackNum == 0:
            readFile = open("test.jpg","rb")
        if currentPackNum == p_num:
            content = readFile.read(512)
            n = len(content)
            print(n)
            p_num += 1
            sendData = struct.pack("!HH512s", 3, p_num,content)
            UDPsocket.sendto(sendData,recvAddr)
            print("发送第%d块数据,大小为%d" % (p_num,n))
            if len(content) < 512:
                finishNum = p_num
        if cmd == 5:
            print("错误应答码：%d"%currentPackNum)
            break
def main():
    global UDPsocket
    UDPsocket = socket(AF_INET,SOCK_DGRAM)
    choose = input("请选择下载或者上传，1下载，2上传：")
    if choose == 1:
        downLoad()
    else:
        upLoad()
    UDPsocket.close()
if __name__ == "__main__":
    main()
